import torch
import jieba
import pandas as pd
import numpy as np
from config import Config
from model import EnhancedModel
from torch.utils.data import DataLoader
from data_processor import TextDataset

class InferenceSystem:
    def __init__(self):
        self.cfg = Config()
        # self._init_jieba()
        self.word2idx = self._build_vocab()
        self.model = self._load_model()

    # def _init_jieba(self):
        # jieba.load_userdict(self.cfg.USER_DICT)
        # 自定义词典内容（需提前创建）
        # 送货速度 10 n
        # 电池续航 10 n
        # 服务态度 10 n

    def _build_vocab(self):
        df = pd.read_csv(self.cfg.DATA_PATH)
        word2idx = {"<PAD>": 0, "<UNK>": 1}
        for text in df["text"]:
            for word in jieba.lcut(text):
                if word not in word2idx:
                    word2idx[word] = len(word2idx)
        return word2idx

    def _load_model(self):
        model = EnhancedModel(len(self.word2idx)).to(self.cfg.DEVICE)
        model.load_state_dict(torch.load(self.cfg.MODEL_PATH, map_location=self.cfg.DEVICE))
        model.eval()
        return model

    def _find_subject_spans(self, probs, threshold=0.6):
        spans = []
        current_start = None
        for i, p in enumerate(probs):
            if p > threshold:
                if current_start is None:
                    current_start = i
                current_end = i
            else:
                if current_start is not None:
                    spans.append((current_start, current_end))
                    current_start = None
        if current_start is not None:
            spans.append((current_start, current_end))
        
        return spans  # Add this return statement
    # 添加评估模型性能的方法
    def evaluate_model(self, test_data_path=None):
        """
        评估模型在测试数据集上的性能
        
        Args:
            test_data_path: 测试数据集路径，如果为None则使用默认数据集
            
        Returns:
            dict: 包含各项评估指标的字典
        """
        # 加载测试数据
        if test_data_path:
            test_df = pd.read_csv(test_data_path)
        else:
            # 使用默认数据集的一部分作为测试集
            df = pd.read_csv(self.cfg.DATA_PATH)
            test_df = df.sample(frac=0.2, random_state=42)  # 随机抽取20%作为测试集
        
        # 创建测试数据集和数据加载器
        test_dataset = TextDataset(test_df, self.word2idx)
        test_loader = DataLoader(
            test_dataset,
            batch_size=self.cfg.BATCH_SIZE,
            shuffle=False
        )
        
        # 初始化评估指标
        subj_correct = 0
        sent_correct = 0
        total_samples = 0
        
        # 在测试集上进行预测
        with torch.no_grad():
            for batch in test_loader:
                inputs = batch["input_ids"].to(self.cfg.DEVICE)
                subj_labels = batch["subject_tags"].to(self.cfg.DEVICE)
                sent_labels = batch["sentiment"].to(self.cfg.DEVICE)
                
                # 模型预测
                subj_pred, sent_pred = self.model(inputs)
                
                # 计算主语标注准确率
                subj_correct += ((torch.sigmoid(subj_pred) > 0.5) == subj_labels.bool()).sum().item()
                
                # 计算情感分类准确率
                sent_correct += (torch.argmax(sent_pred, 1) == sent_labels).sum().item()
                
                total_samples += inputs.size(0)
        
        # 计算最终指标
        subj_acc = subj_correct / (total_samples * self.cfg.MAX_LEN)
        sent_acc = sent_correct / total_samples
        overall_acc = (subj_acc + sent_acc) / 2
        
        # 返回评估结果
        results = {
            "主语标注准确率": f"{subj_acc:.2%}",
            "情感分类准确率": f"{sent_acc:.2%}",
            "综合准确率": f"{overall_acc:.2%}",
            "样本数量": total_samples
        }
        
        return results
    
    # 添加一个简单的方法来打印评估结果
    def print_evaluation(self, test_data_path=None):
        """打印模型评估结果"""
        results = self.evaluate_model(test_data_path)
        print("\n模型评估结果:")
        print("=" * 30)
        for metric, value in results.items():
            print(f"{metric}: {value}")
        print("=" * 30)
        
        return results
    
    def predict(self, text):
        """
        对输入文本进行主语和情感预测
        
        Args:
            text: 输入文本
            
        Returns:
            dict: 包含主语和情感预测结果的字典
        """
        # 分词
        words = jieba.lcut(text)[:self.cfg.MAX_LEN]
        
        # 文本编码
        text_ids = [self.word2idx.get(word, 1) for word in words]  # 未知词用 <UNK> 标记
        pad_len = self.cfg.MAX_LEN - len(text_ids)
        text_ids += [0] * pad_len  # 填充
        
        # 转换为张量并添加批次维度
        inputs = torch.tensor([text_ids], dtype=torch.long).to(self.cfg.DEVICE)
        
        # 模型预测
        with torch.no_grad():
            subj_pred, sent_pred = self.model(inputs)
        
        # 处理主语预测结果
        subj_probs = torch.sigmoid(subj_pred)[0].cpu().numpy()
        
        # 降低阈值以提高主语检测的灵敏度
        spans = self._find_subject_spans(subj_probs, threshold=0.4)
        
        # 提取主语
        subjects = []
        for start, end in spans:
            if start < len(words):  # 确保索引有效
                subject_words = words[start:end+1]
                subjects.append("".join(subject_words))
        
        # 处理情感预测结果
        sent_probs = torch.softmax(sent_pred, dim=1)[0].cpu().numpy()
        sentiment = "正面" if sent_probs[0] > sent_probs[1] else "负面"
        
        # 返回结果
        result = {
            "text": text,
            "subject": subjects[0] if subjects else "未检测到主语",
            "sentiment": sentiment,
            "subject_confidence": float(np.max(subj_probs)) if spans else 0.0,
            "sentiment_confidence": float(np.max(sent_probs))
        }
        
        return result
